# 导入库
import random
import numpy as np
import csv
import time
from RA_Cal.Costfunction import Costfunction_tt

CostF_array = []
class DPGA_new:
    def __init__(self, CF_tt, n_code, population_size, generations, casei,
                 precise_population_ratio=0.1, div_population_ratio=0.2, ):
        # 外部传参
        self.n_code = n_code  # 编码长度
        self.population_size = population_size  # 主群体中个体数量
        self.population_size0 = population_size
        self.generations = generations  # 迭代次数
        self.casei = casei

        # 双种群参数
        self.div_population_size = int(population_size * div_population_ratio)  # 多样性种群大小
        self.migration_rate = 0.06        # 迁移率
        self.migration_interval = 20       # 迁移间隔

        # 内部传参
        self.Init_u = 2 ** (len(CF_tt) - 1)
        self.main_population = np.random.randint(0, 2, (self.population_size * self.Init_u, self.n_code))
        self.div_population = np.random.randint(0, 2, (self.div_population_size, self.n_code))

        self.current_main_diversity = 0  # 当前主种群多样性
        self.current_div_diversity = 0  # 当前多样性种群多样性
        self.number = 0
        self.Tc = 0
        self.deltaa = 0.20

        self.best_fitness = -np.inf  # 初始化最佳适应度为无穷小量
        self.best_solution = None  # 最佳解

        # 修改精确评估相关的成员变量初始化
        self.precise_evaluation_size = int(population_size * precise_population_ratio)  # 精确评估的个体数量
        self.tt_cache = np.array(CF_tt)  # 所有可能用到的TT值

        self.optimize_with_successive_halving()

        self.fitness_fast = CostF_array[0].CostFunction # 粗适应度评估函数
        self.fitness_prec = CostF_array[-1].CostFunction  # 精适应度评估函数
        self.time_c = np.array([1/16, 1/8, 1/4, 1/2, 1], dtype=float)

        self.door_sh = np.array([0.2, 0.15, 0.12, 0.1, 1], dtype=float)
        self.door_sh = np.abs(np.log10(self.tt_cache / self.tt_cache[-1])) * 0.55
        self.door_sh[-1] = 1
        print(self.door_sh)
        self.time_all = 0

    # 改进的迁移策略
    def migrate(self):
        if self.number > 0 and self.number % self.migration_interval == 0:
            num_migrants = int(self.population_size * self.migration_rate)
            main_samp = self.main_population[random.sample(range(self.population_size), num_migrants)]
            dive_samp = self.div_population[random.sample(range(self.div_population_size), num_migrants)]
            temp = main_samp.copy()

            for i in range(num_migrants):
                points = sorted(random.sample(range(1, self.n_code), 2))
                temp[i][points[0]:points[1]] = dive_samp[i][points[0]:points[1]]

            self.main_population = np.vstack((self.main_population, temp))
            self.population_size = len(self.main_population)
        else:
            self.population_size = len(self.main_population)

    def mutation(self):
        # 主种群变异
        self.main_mutation_rate = 0.001
        mutation_mask = np.random.random((self.population_size, self.n_code)) < self.main_mutation_rate
        self.main_population[mutation_mask] = 1 - self.main_population[mutation_mask]

        # 多样性种群变异
        self.div_mutation_rate = 0.005
        mutation_mask = np.random.random((self.div_population_size, self.n_code)) < self.div_mutation_rate
        self.div_population[mutation_mask] = 1 - self.div_population[mutation_mask]

    # 改进的交叉操作
    def crossover(self):
        for i in range(0, self.population_size - 1, 2):
                points = sorted(random.sample(range(1, self.n_code), 2))  # 使用两点交叉
                temp = self.main_population[i].copy()
                self.main_population[i][points[0]:points[1]] = \
                    self.main_population[i + 1][points[0]:points[1]]
                self.main_population[i + 1][points[0]:points[1]] = \
                    temp[points[0]:points[1]]

        for i in range(0, self.div_population_size - 1, 2):
                points = sorted(random.sample(range(1, self.n_code), 3))  # 使用三点交叉
                temp = self.div_population[i].copy()
                for j in range(0, len(points) - 1, 2):
                    self.div_population[i][points[j]:points[j + 1]] = \
                        self.div_population[i + 1][points[j]:points[j + 1]]
                    self.div_population[i + 1][points[j]:points[j + 1]] = \
                        temp[points[j]:points[j + 1]]

    # 计算适应度并保存
    def cal_fitness(self):
        # 向量化计算主种群适应度
        self.y_main_population = np.array([self.fitness_fast(self.casei, ind) for ind in self.main_population])
        # self.y_div_population = np.array([self.fitness_fast(self.casei, ind) for ind in self.div_population])

        # 选择表现最好的个体进行精确评估
        main_indices = np.argsort(self.y_main_population)[-self.precise_evaluation_size:]

        # 精确评估
        precise_fitness_main = np.zeros(len(self.main_population)) - np.inf
        if self.Tc < len(self.tt_cache) - 1:
            precise_fitness_main[main_indices] = np.array([
                self.fitness_prec(self.casei, self.main_population[i]) for i in main_indices])
        else:
            precise_fitness_main[main_indices] = self.y_main_population[main_indices]

        if self.number % 1 == 0:
            kjj = np.max(np.abs((precise_fitness_main[main_indices] - self.y_main_population[main_indices])))
            # print(kjj)
            if kjj > self.door_sh[self.Tc]:
                self.Tc = min(self.Tc + 1, len(self.tt_cache) - 1)

        self.y_main_population[main_indices] = precise_fitness_main[main_indices]

        # 更新全局最优
        aa = np.max(precise_fitness_main)
        if aa > self.best_fitness:
            self.best_fitness = aa
            self.best_code = self.main_population[np.argmax(precise_fitness_main)]

            # Curve_result_file = "DPGAnew_result_" + str(self.casei) + ".csv"
            # BestCode_result_file = "DPGAnew_code_" + str(self.casei) + ".csv"
            # with open(Curve_result_file, 'a', encoding='utf-8', newline='') as f:
            #     write = csv.writer(f)
            #     write.writerow([self.number, self.best_fitness.round(4)])
            # with open(BestCode_result_file, 'w', encoding='utf-8', newline='') as f:
            #     write = csv.writer(f)
            #     write.writerow(list(self.best_code))

    def selection(self):
        if self.number > 0 and self.number % self.migration_interval == 0:
            main_indices = np.argsort(self.y_main_population)[-self.population_size0 :]
            self.main_population = self.main_population[main_indices]
        else:
            competitors = np.random.randint(0, self.population_size,
                                            (self.population_size, 3))
            winners = np.argmax(self.y_main_population[competitors], axis=1)
            self.main_population = self.main_population[competitors[np.arange(len(winners)), winners]]

        # 多样性种群选择
        competitors = np.random.randint(0, self.div_population_size,
                                        (self.div_population_size, 3))
        winners = np.argmax(self.diversity_scores[competitors], axis=1)
        self.div_population = self.div_population[competitors[np.arange(len(winners)), winners]]

    def optimize_with_successive_halving(self):
        # 使用Successive Halving方法优化主种群
        self.main_population, self.y_main_population = successive_halving(
            self.main_population, casei=self.casei, tt_cache=self.tt_cache)

        self.best_fitness = np.max(self.y_main_population)
        self.best_code = self.main_population[np.argmax(self.y_main_population)]

    def cal_diver_scores(self):
        self.aver_code_main = np.sum(self.main_population, axis=0) / self.population_size - 0.001
        self.aver_code_main = np.array(self.aver_code_main * 2, dtype=int)

        self.main_diversity = np.zeros(self.population_size)
        self.div_diversity = np.zeros(self.div_population_size)
        self.main_diversity = np.array(
            [np.sum(ind ^ self.aver_code_main) for ind in self.main_population]) / self.n_code
        self.div_diversity = np.array(
            [np.sum(ind ^ self.aver_code_main) for ind in self.div_population]) / self.n_code

        self.diversity_scores = np.zeros(self.div_population_size)
        self.diversity_scores[self.div_diversity < self.deltaa] = 1 - (
                self.deltaa - self.div_diversity[self.div_diversity < self.deltaa]) / self.deltaa
        self.diversity_scores[self.div_diversity > self.deltaa] = 1 - (
                self.div_diversity[self.div_diversity > self.deltaa] - self.deltaa) / (1 - self.deltaa)

        self.current_main_diversity = np.median(self.main_diversity)
        self.current_div_diversity = np.median(self.div_diversity)

    def run(self):
        for self.number in range(self.generations):
            sti = time.time()

            self.migrate()
            self.crossover()
            self.mutation()

            self.cal_fitness()
            self.cal_diver_scores()
            self.selection()

            self.fitness_fast = CostF_array[self.Tc].CostFunction
            self.time_all += self.population_size * self.time_c[self.Tc]
            if self.Tc < len(self.tt_cache) - 1:
                self.time_all += self.precise_evaluation_size

            self.deltaa = min(self.current_main_diversity * 5, 0.5)
            Curve_result_file = "CURVE_DPGA_adpt_" + str(self.casei) + ".csv"
            with open(Curve_result_file, 'a', encoding='utf-8', newline='') as f:
                write = csv.writer(f)
                write.writerow([self.time_all, self.best_fitness.round(4),
                                self.number, round(self.deltaa, 5), self.tt_cache[self.Tc]])

            if self.number % 100 == 0:
                print(f"num {self.number}: "
                    f"t {round(self.time_all, 0)}: "
                    f"Best Fitness={self.best_fitness:.5f}, "
                    f"Delta={self.deltaa:.5f}, "
                    f"Main Diversity={self.current_main_diversity:.5f}, "
                    f"Div Diversity={self.current_div_diversity:.5f}, "
                    f"Tc={self.Tc}, "
                    f"Time={time.time() - sti:.2f}s")

        # 获取最终最优解
        all_solutions = self.main_population
        all_fitness = self.y_main_population
        best_idx = np.argmax(all_fitness)
        best_code = all_solutions[best_idx]

        # 保存结果
        np.set_printoptions(threshold=np.inf)
        return best_code, self.best_fitness

def successive_halving(initial_population, casei, tt_cache):
    print(f"Successive Halving使用的TT值: {tt_cache}")

    for tti in tt_cache:
        CostF_array.append(Costfunction_tt(tti))

    population = initial_population.copy()
    current_tt = None  # 追踪当前TT值

    for step, TT_value in enumerate(tt_cache):
        print(f"步骤 {step + 1}/{len(tt_cache)}: 使用TT={TT_value}评估{len(population)}个个体")
        start_time = time.time()  # 记录步骤开始时间

        # 向量化计算适应度
        fitness_values = np.array([
            CostF_array[step].CostFunction(casei, ind) for ind in population
        ])
        fitness_values[fitness_values == None] = 0.0  # 处理None值

        end_time = time.time()  # 记录步骤结束时间
        elapsed_time = end_time - start_time  # 计算消耗时间
        print(f"步骤 {step + 1} 消耗时间: {elapsed_time:.2f} 秒")
        # 如果是最后一步，返回所有个体
        if step == len(tt_cache) - 1:
            break

        # 否则，保留10%表现最好的个体
        half_size = max(int(len(population) * 0.5), 1)
        sorted_indices = np.argsort(fitness_values)
        population = np.array([population[i] for i in sorted_indices[-half_size:]])  # 确保返回的是NumPy数组

        # # 否则，保留一半表现最好的个体
        # sorted_indices = np.argsort(fitness_values)
        # half_size = max(len(population) // 2, 1)
        # population = np.array([population[i] for i in sorted_indices[:half_size]])  # 确保返回的是NumPy数组

    # 返回最终种群和它们的适应度值
    return population, np.array(fitness_values)  # 确保返回的是NumPy数组
